# -*- coding: utf-8 -*-
"""
轮动策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
from Functions import *
from Evaluate import *
from Config import *
from joblib import Parallel, delayed
import sys
import warnings

warnings.filterwarnings('ignore')
pd.set_option('display.max_rows', 1000)
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行

if __name__ == '__main__':
    # 动态读取config里面配置的shift_name脚本
    Shift = __import__('shift.%s' % shift_name, fromlist=('',))

    # =====取出Shift中配置的变量
    # 从Shift中取出 因子列表factor_list
    factor_list = Shift.factor_list
    # 从Shift中取出 回测持仓周期hold_period
    hold_period = Shift.hold_period
    # 从Shift中取出 offset
    offset = Shift.offset
    # 从Shift中取出 选策略数量select_num
    select_num = Shift.select_num
    # 从Shift中取出 资金曲线文件all_equity_file
    all_equity_file = Shift.all_equity_file
    # 从Shift中取出 子策略池base_strategy_dict
    base_strategy_dict = Shift.base_strategy_dict
    # 从Shift中取出 策略名stg_name
    stg_name = Shift.stg_name

    # =手动运行标识。手动运行的时候才会出画图，脚本调用直接跳过
    if_manual = True  # 默认手工运行
    if len(sys.argv) > 1:
        if_manual = False
        if_display_stat = False

        # 取出遍历的offset
        offset = int(sys.argv[1])

        if len(sys.argv) > 2:
            factor_col = sys.argv[2]  # 取出遍历因子列
            factor = factor_col.split('_')[0]  # 取出对应的因子
            param = int(factor_col.split('_')[1])  # 取出对应的参数

            if_ascending = eval(sys.argv[3])  # 取出排序方式

            factor_list = [(factor, if_ascending, param, 1)]  # 根据遍历因子列得到因子列表

            select_num = int(sys.argv[4])  # 取出选策略数量
            _save_path = sys.argv[5]  # 取出保存路径

        all_equity_file = get_equity_file(base_strategy_dict, base_config_path, hold_period, offset)

    print('开始回测...')
    print('持仓周期：', hold_period)
    print('offset：', offset)
    print('因子列表：', factor_list)
    print('选策略数量：', select_num)

    # 将配置好的因子列表解析为 因子列表和参数列表
    factor_names = []
    factor_params = []
    factor_cols = []
    for factor_info in factor_list:
        factor_names.append(factor_info[0])
        factor_params.append(factor_info[2])
        factor_cols.append(factor_info[0] + '_' + str(factor_info[2]))

    # =====读取数据
    # 并行或串行，依次取出每个资金曲线的数据，最终生成一个总df，每个资金曲线当作币种进行处理
    multiply_process = False  # 是否并行。
    if multiply_process:
        equity_list = Parallel(n_jobs=n_jobs)(delayed(read_equity)(path, base_strategy_dict, hold_period, offset, benchmark) for path in all_equity_file)
    else:
        equity_list = []
        for path in all_equity_file:
            _equity = read_equity(path, base_strategy_dict, hold_period, offset, benchmark)
            equity_list.append(_equity)

    # 合并并整理数据
    base_equity = pd.concat(equity_list, axis=0)
    base_equity = base_equity.drop_duplicates(subset=['candle_begin_time', 'symbol'], keep='last')
    base_equity = base_equity.sort_values(['candle_begin_time', 'symbol']).reset_index(drop=True)  # 排序并重置索引

    # =====计算因子
    # 并行或串行，依次取出每个资金曲线的数据，计算因子
    multiply_process = False  # 是否并行。在测试的时候可以改成False，实际跑的时候改成True
    if multiply_process:
        df_list = Parallel(n_jobs=n_jobs)(delayed(calc_factors)(base_equity, symbol, factor_names, factor_params) for symbol in base_equity['symbol'].unique())
    else:
        df_list = []
        for symbol in base_equity['symbol'].unique():
            _equity = calc_factors(base_equity, symbol, factor_names, factor_params)
            df_list.append(_equity)
    # 合并并整理数据
    df = pd.concat(df_list, axis=0)
    if df[factor_cols].isnull().sum().max() == len(df):
        nan_df = df[factor_cols].isnull().sum()
        nan_df = nan_df[nan_df == nan_df.max()]
        print(f'{list(nan_df.index)}因子全部为空，请检查因子')
        exit()
    df = df.dropna(subset=factor_cols, axis=0)  # 删除因子值为空的数据
    df = df[(df['candle_begin_time'] >= start_date) & (df['candle_begin_time'] <= end_date)]  # 根据回测开始时间、结束时间截取数据
    if df.empty:
        print('日期配置有误，导致资金曲线为空，请检查日期及资金曲线文件')
        exit()
    df = df.sort_values(['candle_begin_time', 'symbol']).reset_index(drop=True)  # 排序

    # 保存一下轮动策略的中间数据用于轮动模拟器
    save_path = os.path.join(back_test_path, f'{stg_name}_轮动策略中间数据_{hold_period}_{offset}.pkl')
    save_inter_data(df, factor_cols, save_path)

    # ===根据因子选策略
    df = select_strategy(df, factor_list, select_num)

    # ===计算资金曲线
    df = cal_rot_net_value(df, hold_period, spot_buy_crate, spot_sell_crate, swap_buy_crate, swap_sell_crate)

    # 判断是否展示统计信息
    if if_display_stat:
        # ===根据计算出的资金曲线数据得到轮动的评价指标
        display_stat_data(df, info_path, stg_name, hold_period, offset)

    # ===计算资金曲线
    # =创造空的时间周期表，用于填充不选币的周期
    _benchmark = get_benchmark(df['candle_begin_time'].min(), df['candle_begin_time'].max(), freq='1H')
    equity = pd.merge(_benchmark, df[['candle_begin_time', '策略名称', '选币', '策略数量', 'label']], on='candle_begin_time', how='left')
    equity[['策略名称', '选币', '策略数量', 'label']] = equity[['策略名称', '选币', '策略数量', 'label']].fillna(method='ffill')
    equity['轮动涨跌幅'] = revise_data_length(df['每小时涨跌幅'].sum(), len(equity))
    equity['_'] = equity['轮动涨跌幅'].fillna(method='bfill')
    if equity['_'].isnull().sum() > 0:
        index = equity[equity['_'].isnull()].index[0]
        equity = equity[equity.index < index].reset_index(drop=True)
    equity.drop('_', axis=1, inplace=True)
    equity['轮动资金曲线'] = (equity['轮动涨跌幅'] + 1).cumprod()

    # 保存资金曲线文件
    save_path = os.path.join(back_test_path, f'{stg_name}_轮动资金曲线_{hold_period}_{offset}.csv')
    equity.to_csv(save_path, encoding='gbk', index=False)

    # ===计算策略评价指标
    rtn, year_return, month_return = strategy_evaluate(equity, net_col='轮动资金曲线', pct_col='轮动涨跌幅')
    print(rtn, '\n\n分年收益率：\n', year_return, '\n\n分月收益率：\n', month_return)

    # ===画图：
    if if_manual:
        # 画资金曲线图1
        plot_ratation_net_value1(equity, factor_list, rtn)
        if select_num == 1:  # 当选策略数量为1时，画轮动资金曲线2
            plot_ratation_net_value2(equity, base_equity, factor_list, rtn)
    else:
        if len(sys.argv) > 2:
            # 记录遍历回测的结果
            res = pd.DataFrame()
            res.loc[0, '因子'] = [factor_list]
            res.loc[0, '持仓周期'] = hold_period
            res.loc[0, 'offset'] = offset
            res.loc[0, '选策略数量'] = select_num
            res.loc[0, '策略名'] = stg_name
            res.loc[0, '换仓次数'] = equity['label'].max()
            year_return['涨跌幅'] = year_return['涨跌幅'].apply(lambda x: float(x[:-1]))
            res['各年收益'] = [year_return.values.reshape(-1).tolist()]
            # 记录回测的各项指标
            stg_res = pd.concat([res, rtn.T], axis=1)
            # 保存文件
            if os.path.exists(_save_path):  # 如果文件存在，往原有的文件中添加新的结果
                stg_res.to_csv(_save_path, encoding='gbk', index=False, header=False, mode='a')
            else:  # 如果文不件存在，常规的to_csv操作
                stg_res.to_csv(_save_path, encoding='gbk', index=False)
